package com.stegniy.binaryTree;

import java.util.*;

public class BinaryTree<T extends Comparable<T>> {
    Node<T> root;

    public BinaryTree(T element) {
        root = new Node<T>(element,null,null);
    }

    public boolean contains(T element) {
        return contains(element, root);
    }

    private boolean contains(T element, Node<T> node) {
        if (node == null) {
            return false;
        } else {
            if (node.data.compareTo(element) == 0)
                return true;
            else if (node.data.compareTo(element) > 0)
                return contains(element, node.leftChild);
            else {
                return contains(element, node.rightChild);
            }
        }
    }

    public int size() {
        // size = 1 if tree contains only root node
        return size(root) + 1;
    }
    private int size(Node node) {
        int result = 0;
        if(node.leftChild != null)
            result = result + size(node.leftChild) + 1;
        if(node.rightChild != null)
            result = result + size(node.rightChild) + 1;
        return result;
    }

    public void add(T element) {
        Node<T> node = new Node<T>(element,null,null);
        insertNode(root,node);
    }

    private Node<T> insertNode(Node<T> subTree, Node<T> newNode) {
        if (subTree == null) {
            subTree = newNode;
        } else if ((subTree.data).compareTo(newNode.data) < 0){
            subTree.rightChild = insertNode(subTree.rightChild, newNode);
        } else if ((subTree.data).compareTo(newNode.data) > 0){
            subTree.leftChild = insertNode(subTree.leftChild, newNode);
        } else {
            throw new IllegalArgumentException("Element already exist " + newNode.data);
        }
        return subTree;
    }

    public void deleteAll() {
        root = null;
    }

    public void delete(T element) {
        this.deleteNode(root,element);
    }

    private Node<T> deleteNode(Node<T> subtree,  T element) {
        if (subtree != null) {
            if ((subtree.data).compareTo(element) < 0) {
                subtree.rightChild = deleteNode(subtree.rightChild, element);
            } else if ((subtree.data).compareTo(element) > 0) {
                subtree.leftChild = deleteNode(subtree.leftChild, element);
            } else {
                if ((subtree.leftChild != null) && (subtree.rightChild != null)) {
                    Node<T> node = findLeftmostChild(subtree.rightChild);
                    subtree.data = node.data;
                    subtree.rightChild = deleteNode(subtree.rightChild,node.data);
                } else if (subtree.leftChild != null) {
                    subtree = subtree.leftChild;
                } else {
                    subtree = subtree.rightChild;
                }
            }

        } else {
            throw new IllegalArgumentException("No such element");
        }
        return subtree;
    }

    private Node<T> findLeftmostChild(Node<T> subtree){
        assert (subtree != null);
        while (subtree.leftChild != null) {
            subtree = subtree.leftChild;
        }
        return subtree;
    }

    public List<T> depthFirstTraversal() {
        List<T> nodeList = new ArrayList<T>();
        Stack<Node<T>> nodeStack = new Stack<Node<T>>();
        nodeStack.push(root);
        while (!nodeStack.isEmpty()){
            Node<T> node = nodeStack.pop();
            nodeList.add(node.data);
            if (node.rightChild != null) {
                nodeStack.push(node.rightChild);
            }
            if (node.leftChild != null) {
                nodeStack.push(node.leftChild);
            }
        }
        return nodeList;
    }

    public List<T> breadthFirstTraversal() {
        List<T> nodeList = new ArrayList<T>();
        Queue<Node<T>> nodeQueue = new LinkedList<Node<T>>();
        nodeQueue.add(root);
        while (!nodeQueue.isEmpty()) {
            Node<T> node = nodeQueue.poll();
            nodeList.add(node.data);
            if (node.leftChild != null) {
                nodeQueue.add(node.leftChild);
            }
            if (node.rightChild != null) {
                nodeQueue.add(node.rightChild);
            }
        }
        return nodeList;
    }

    public Node<T> findNode(T element) {
        if (element == null) return null;
        Node<T> current = root;
        while ((current.data).compareTo(element) != 0) {
            if ((current.data).compareTo(element) > 0) {
                current = current.leftChild;
            } else if ((current.data).compareTo(element) < 0) {
                current = current.rightChild;
            }
            if (current == null) return null;
        }
        return current;

    }
}

class Node<T extends Comparable<T>>{
    T data;
    Node<T> leftChild;
    Node<T> rightChild;

    Node(T element, Node<T> left, Node<T> right){
        this.data = element;
        leftChild = left;
        rightChild = right;
    }
}
